#include "stdio.h"
#include "string.h"
#include "lua/lua.h"
#include "srtp2/srtp.h"

static int init = 0;
static srtp_t ctx1 = NULL;
static srtp_t ctx2 = NULL;
typedef srtp_err_status_t (*unprotect)(srtp_t ctx, void *srtp_hdr, int *len_ptr);

int module_init(lua_State *L){
    if (init) return 0;

    srtp_init();
    return 0;
}

int module_deinit(lua_State *L){
    if(!init) return 0;

    srtp_shutdown();
    return 0;
}

int destroy(lua_State *L){
    if (ctx2) {
        srtp_dealloc(ctx2);
        ctx2 = NULL;
    }

    if (ctx1) {
        srtp_dealloc(ctx1);
        ctx1 = NULL;
    }

	return 0;
}

int set_key(lua_State *L){
    if (ctx1 || ctx2) destroy(L);

    srtp_policy_t policy;

    memset(&policy, 0,sizeof(policy));
    srtp_crypto_policy_set_aes_cm_128_hmac_sha1_80(&policy.rtp);
    srtp_crypto_policy_set_aes_cm_128_hmac_sha1_80(&policy.rtcp);

    policy.ssrc.value = 0;
    policy.window_size = 8192;
    policy.allow_repeat_tx = 1;
    policy.next = NULL;
    srtp_err_status_t r0 = srtp_err_status_ok;

    //init send context
    policy.ssrc.type = ssrc_any_outbound;
    policy.key = lua_tostring(L, 1);  
    if ((r0 = srtp_create(&ctx1, &policy)) != srtp_err_status_ok) {
        return 0;
    }

    // init recv context
    policy.ssrc.type = ssrc_any_inbound;
    policy.key = lua_tostring(L, 2);  
    if ((r0 = srtp_create(&ctx2, &policy)) != srtp_err_status_ok) {
        return 0;
    }
    return 0;
}

int decrypt(lua_State *L){
    char buff[1500];

    size_t len;
    char *data = lua_tolstring(L, 1, &len);
    memcpy(buff, data, len);

    unprotect f = NULL;
    unsigned char p1 = data[1];
    f = p1 >= 192 && p1 <= 223 ? srtp_unprotect_rtcp : srtp_unprotect;

    int r_len = len;
    srtp_err_status_t ret = f(ctx1, buff, &r_len);
    if (ret != srtp_err_status_ok) ret = f(ctx2, buff, &r_len);
    if (ret != srtp_err_status_ok) return 0;

    lua_pushlstring(L, buff, r_len);
    return 1;
}